from zkl_pyutils_fsspec import FsLike, resolve_fs, resolve_parent_fs

from zkl_aiutils_datasets.basics import NamedDatasets
from zkl_aiutils_datasets.scripting import load_dataset_from_metadata, load_dataset_from_script


def load_dataset_by_scripting(fs: FsLike, *args, **kwargs):
    fs = resolve_fs(fs)
    exceptions = []

    try:
        return load_dataset_by_scripting_from_itself(fs, *args, **kwargs)
    except Exception as exception:
        exceptions.append(exception)

    try:
        return load_dataset_by_scripting_from_ancestors(fs, *args, **kwargs)
    except Exception as exception:
        exceptions.append(exception)

    raise ExceptionGroup(f"Failed to load dataset from {fs}", exceptions)


def load_dataset_by_scripting_from_ancestors(fs: FsLike, *args, **kwargs):
    parent_fs, child_name = resolve_parent_fs(fs)
    parent_dataset = load_dataset_by_scripting(parent_fs, *args, **kwargs)
    if not isinstance(parent_dataset, NamedDatasets):
        raise TypeError(f"Dataset at {parent_fs} is not a NamedDatasets.")
    return parent_dataset.named_children[child_name]


def load_dataset_by_scripting_from_itself(fs: FsLike, *args, **kwargs):
    fs = resolve_fs(fs)
    exceptions = []

    try:
        return load_dataset_from_script(fs, *args, **kwargs)
    except Exception as exception:
        exceptions.append(exception)

    try:
        return load_dataset_from_metadata(fs, *args, **kwargs)
    except Exception as exception:
        exceptions.append(exception)

    raise ExceptionGroup(f"Failed to load dataset from {fs}", exceptions)
